package com.shuanwang.nl2sql_demo;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.concurrent.TimeUnit;

public class Nl2SqlExample {
    /*public static GenerationResult callWithMessage() throws ApiException, NoApiKeyException, InputRequiredException {
        // 创建 Generation 实例
        Generation gen = new Generation();

        // 系统消息：定义助手的角色
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你是一个专业的自然语言转 SQL 工具，能够根据数据库模式和问题生成准确的 SQL 查询。")
                .build();

        // 用户消息：包含数据库模式和问题
        String prompt = """
                请根据以下数据库模式和问题生成 SQL 查询：
                
                数据库模式:
                表名: employees
                字段: id (整数), name (字符串), age (整数), department (字符串)
                
                问题: 查询所有年龄大于 30 岁的员工姓名。
                """;
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();

        // 构建调用参数
        GenerationParam param = GenerationParam.builder()
                // 替换为您的百炼 API Key
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                // 指定模型名称
                .model("qwen-plus") // 或 "qwen-max"，根据需求选择
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();

        // 调用 API
        return gen.call(param);
    }

    public static void main(String[] args) {
        try {
            // 调用 API 并获取结果
            GenerationResult result = callWithMessage();
            // 输出生成的 SQL 查询
            System.out.println("生成的 SQL 查询: " + result.getOutput().getChoices().get(0).getMessage().getContent());
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 错误处理
            System.err.println("错误信息：" + e.getMessage());
            System.out.println("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code");
        }
        System.exit(0);
    }*/

    public static GenerationResult callWithMessage() throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("You are a helpful assistant.")
                .build();
        String schema;
        try {
            schema = new String(Files.readAllBytes(Paths.get("Schema.json")));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        //String prompt = schema + "熟悉schema,将开始向你提问";
        String str = "请生成标准的 SQL 查询，避免使用非标准语法或复杂函数。只输出SQL语句可以直接在数据库中查询。";
        String prompt = str + """
                查询每个测站的最新水位 (RZ) 和流量 (Q)，并按时间戳 (TM) 排序。
                """;
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey("sk-69f2b0634b304d49abafb1df188d4440")
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model("qwen-max")
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.call(param);
    }

    public static void main1(String[] args) throws IOException {
        try {
            GenerationResult result = callWithMessage();
            System.out.println(result.getOutput().getChoices().get(0).getMessage().getContent());
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.err.println("错误信息：" + e.getMessage());
            System.out.println("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code");
        }
        System.exit(0);
    }


    private static boolean flag = false;
    private static volatile int i = 0;
    public static void main(String[] args) {
        new Thread(() -> {
            try {
                TimeUnit.MILLISECONDS.sleep(100);
                flag = true;
                System.out.println("flag 被修改成 true");
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }).start();

        while (!flag) {
            i++;
        }

        System.out.println("程序结束,i=" + i);
    }

}
